import os
import numpy as np
from read_obj_init_pose import read_obj_pose


def get_mat_from_pose(pose):
    mat = np.eye(4)
    mat[:3, :3] = pose["rotation"]
    mat[:3, 3:] = pose["translation"]
    return mat


def get_pose_from_mat(mat):
    pose = {
        "rotation": mat[:3, :3],
        "translation": mat[:3, 3:],
    }
    return pose


def read_obj_poses(video_dir):
    """
    input: video_dir
    return:
    poses: len = 300
    """

    poses = []

    init_pose = read_obj_pose(os.path.join(video_dir, "obj_init_pose", "00000.pkl"))
    assert "scale" in init_pose
    poses.append(init_pose)
    init_pose_mat = get_mat_from_pose(init_pose)

    for i in range(1, 300):
        delta_pose = read_obj_pose(os.path.join(video_dir, "obj_delta_pose", str(i).zfill(5) + ".pkl"))
        delta_pose_mat = get_mat_from_pose(delta_pose)
        current_pose_mat = delta_pose_mat @ init_pose_mat
        current_pose = get_pose_from_mat(current_pose_mat)
        current_pose["scale"] = init_pose["scale"]
        poses.append(current_pose)
    
    return poses


if __name__ == "__main__":
    video_dir = ""  # TODO
    pose = read_obj_poses(video_dir)
